import numpy as np
import pandas as pd
from joblib import dump
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler


def shrink(filename):
    try:
        df = pd.read_csv(filename, nrows=1000000, low_memory=False)
        new_filename = 'all_small.csv'
        df.to_csv(new_filename, index=False)
        print(f"Saved file rows to: {new_filename}")
    except Exception as e:
        print(f"Error processing {filename}: {e}")
    return new_filename


def preprocess(filename):
    try:
        # 读取CSV文件，修正列名中的空格
        df = pd.read_csv(filename, dtype={85: str})
        df.columns = df.columns.str.strip()  # 去除列名中的前后空格

        # 指定使用的列，确保列名正确
        use_columns = [
            'Source IP', 'Source Port', 'Destination IP', 'Destination Port',
            'Protocol', 'Flow Duration', 'Total Fwd Packets',
            'Total Length of Fwd Packets', 'Fwd Packet Length Max',
            'Fwd Packet Length Min', 'Flow Bytes/s', 'Flow Packets/s',
            'Flow IAT Mean', 'Flow IAT Std', 'Flow IAT Max', 'Label'
        ]

        # 确保所有指定的列都存在于数据中
        missing_columns = [col for col in use_columns if col not in df.columns]
        if missing_columns:
            raise ValueError(f"Missing columns in the data: {missing_columns}")

        df = df[use_columns]  # 只使用指定的列
        # 替换无限值并删除缺失值
        df.replace([np.inf, -np.inf], pd.NA, inplace=True)
        df.dropna(inplace=True)

        # 特征工程
        # 标签编码，假设标签列名为 'Label'
        df['Label'] = df['Label'].astype('category').cat.codes

        # 排除非数值列
        numeric_cols = df.columns[(df.dtypes == 'float64') | (df.dtypes == 'int64')]

        # 分割数据为训练集和测试集 (70% 训练, 30% 测试)
        X = df[numeric_cols]
        y = df['Label']

        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

        # 特征缩放
        scaler = StandardScaler()
        X_train_scaled = scaler.fit_transform(X_train)
        X_test_scaled = scaler.transform(X_test)

        # 构建DataFrame回来用于保存
        train_df = pd.DataFrame(X_train_scaled, columns=X_train.columns)
        test_df = pd.DataFrame(X_test_scaled, columns=X_test.columns)
        train_df['Label'] = y_train.values
        test_df['Label'] = y_test.values

        # 保存处理后的训练集和测试集
        train_filename = 'train_mini.csv'
        test_filename = 'test_mini.csv'
        train_df.to_csv(train_filename, index=False)
        test_df.to_csv(test_filename, index=False)

        # 保存scaler到文件
        scaler_filename = '../../model/scaler_mini.joblib'
        dump(scaler, scaler_filename)
        print(f"Scaler saved to {scaler_filename}")

        print(f"Processed and saved: {train_filename} and {test_filename}")
    except Exception as e:
        print(f"Error processing {filename}: {e}")


if __name__ == '__main__':
    # small_file = shrink('all_data.csv')
    preprocess('all_small.csv')
    # preprocess('all_small.csv')
